package geppetto.phraseHMM;


import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.common.StaticTools;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.SparseTranslationTable;

import java.io.IOException;
import java.io.PrintStream;



/** class to train and save an (HMM) alignment model */
public class NewSaveModel {

	public static void main(String[] args) throws IOException {
		String corpusFile = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int numberIterations = Integer.parseInt(args[3]); // 5
		boolean trainWithResults = Boolean.parseBoolean(args[4]);
		int numberIterationsWithResults = Integer.parseInt(args[5]);
		double smooth = Double.parseDouble(args[6]);
		String dir = args[7];
		String modelName = args[8];
		
		System.out.println("Saving Models experiment: ");
		System.out.println("Corpus " + corpusFile);
		System.out.println("Size " + size);
		System.out.println("Smoothing " + smooth);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of EM iterations " + numberIterations);
		System.out.println("Train with results " + trainWithResults + " iter "+ numberIterationsWithResults);
		System.out.println("OutputDir " + dir);
		System.out.println("Model Type " + modelName);
		System.out.println("---");
		BilingualCorpus corpusF = BilingualCorpus.getCorpusFromFileDescription(
				corpusFile, size, maxSentenceSize);
		BilingualCorpus corpusB = corpusF.reverse();
		String baseDir = dir + "/" + modelName + "/" + corpusF.getName() + "/" + size;
		String modelDir = baseDir + "/model/";
		System.out.println("Savaing models to "+ modelDir);
		
		
		//Model 1 code
		IBMM1 m1F;
		IBMM1 m1B;
		m1F = new IBMM1(corpusF,smooth);
		m1F.train(numberIterations,false,"");
		m1B = new IBMM1(corpusB,smooth);
		m1B.train(numberIterations,false,"");
		String M1Dir = modelDir + "/M1/";
		StaticTools.createDir(M1Dir);
		m1F.saveModel(M1Dir + "forward");
		m1B.saveModel(M1Dir + "backward");

		// NOTE: we're intentionally clobbering the translation table of
		// model1! (since we aren't using it anymore)
		RegularHMM mhmmF = null;
		RegularHMM mhmmB = null;
		
		String MHMMDir = modelDir + "MHMM/";
		if (modelName.equalsIgnoreCase("regularHMM")) {
			

			mhmmF = new RegularHMM(corpusF, m1F._tb,smooth);
			if(trainWithResults){
				mhmmF.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir+ "forward/");
			}else{
				mhmmF.train(numberIterations, true, MHMMDir+ "forward/");
			}
		
			mhmmB = new RegularHMM(corpusB, m1B._tb,smooth);
			if(trainWithResults){
				mhmmB.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir + "backward/");
			}else{
				mhmmB.train(numberIterations, true, MHMMDir + "backward/");
			}

			StaticTools.createDir(MHMMDir);
			mhmmF.saveModel(MHMMDir + "forward/");
			mhmmB.saveModel(MHMMDir + "backward/");
		} else{ 
			//MUST READ PROJECTION SPECIFIC PARAMETERS 
			double slack = Double.parseDouble(args[9]); 
			double epsilon = Double.parseDouble(args[10]); 
			int maxStepSize = Integer.parseInt(args[11]);
			int maxNumberOfProjectionIterations = Integer.parseInt(args[12]);	
			System.out.println("Slack " + slack);
			System.out.println("Epsilon " + epsilon);
			System.out.println("Max Number of Steps " + maxStepSize);
			System.out.println("Max Number of projection iter " + maxNumberOfProjectionIterations);
			System.out.println("Slack " + slack);
			if (modelName.equalsIgnoreCase("symmetric")) {
				SymmetricHMM mhmm;						
				mhmm = new SymmetricHMM(corpusF, corpusB, m1F._tb,m1B._tb,smooth,epsilon,slack,maxStepSize,maxNumberOfProjectionIterations);
				if(trainWithResults){
					mhmm.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir);
				}else{
					mhmm.train(numberIterations, true, MHMMDir);
				}
					
				StaticTools.createDir(MHMMDir);
				mhmm.saveModel(MHMMDir);
				mhmmF = mhmm.forward;
				mhmmB = mhmm.backward;
			} else if (modelName.equalsIgnoreCase("bijective")) {				
				mhmmF = new BijectiveHMM(corpusF, m1F._tb, smooth,epsilon,slack,maxStepSize,maxNumberOfProjectionIterations);
				if(trainWithResults){
					mhmmF.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir+ "forward/");
				}else{
					mhmmF.train(numberIterations, true, MHMMDir+ "forward/");
				}
				
				mhmmB = new BijectiveHMM(corpusB, m1B._tb,smooth,epsilon,slack,maxStepSize,maxNumberOfProjectionIterations);
				if(trainWithResults){
					mhmmB.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir+"backward/");
				}else{
					mhmmB.train(numberIterations, true, MHMMDir+"backward/");
				}
				StaticTools.createDir(MHMMDir);
				mhmmF.saveModel(MHMMDir + "forward/");
				mhmmB.saveModel(MHMMDir + "backward/");
			}else if (modelName.equalsIgnoreCase("symmetric-withBijective")) {
				
				String bijectiveDir = dir + "bijective" + "/" + corpusF.getName() + "/" + size +"/model/MHMM/";
		
				SparseTranslationTable ft = SparseTranslationTable.LoadTranslationTable(corpusF, bijectiveDir+"forward/");
				SparseTranslationTable bt = SparseTranslationTable.LoadTranslationTable(corpusF, bijectiveDir+"backward/");		
				SymmetricHMM mhmm;						
				mhmm = new SymmetricHMM(corpusF, corpusB, ft,bt,smooth,epsilon,slack,maxStepSize,maxNumberOfProjectionIterations);
				if(trainWithResults){
					mhmm.trainWithResults(numberIterationsWithResults, BilingualCorpus.DEV_CORPUS, true, MHMMDir);
				}else{
					mhmm.train(numberIterations, true, MHMMDir);
				}
				StaticTools.createDir(MHMMDir);
				mhmm.saveModel(MHMMDir);
				mhmmF = mhmm.forward;
				mhmmB = mhmm.backward;
			}else {
				System.out.println("Unknown Model" + modelName);
				System.exit(1);
			}
		}		
			
		System.out.println("Done with the SAVE MODEL Experience");
	}

}
